﻿using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Threading.Tasks;
using System.Windows.Forms;
using Ast.Core.Pe;
using Ast.Core.Pe.Definitions;
using Ast.Core.Services;
using Ast.Core.Services.Extensions;
using Ast.Framework.Forms;
using Castle.Core.Logging;

namespace Ast
{
    [FormMetadata(FormType.Pe, IsTabForm = true)]
    public partial class PeForm : Form
    {
        private readonly IPeService _peService;
        private readonly IDynamicValueAccessor _dynamicValueAccessor;

        private static readonly string[] PeDirectories = {"Export table", "Import table", "Resource table", "Exception table", "Certificate table", "Base relocation", "Debugging information", "Architecture-specific data", "Global pointer register", "Thread local storage (TLS)", "Load configuration", "Bound import table", "Import address table", "Delay import descriptor", "The CLR header", "Reserved"};

        public PeForm(IPeService peService, IDynamicValueAccessor dynamicValueAccessor)
        {
            _peService = peService;
            _dynamicValueAccessor = dynamicValueAccessor;

            Logger = NullLogger.Instance;

            InitializeComponent();
        }

        public ILogger Logger { get; set; }

        private async void btnPe_Click(object sender, EventArgs e)
        {
            if (openFileDialog1.ShowDialog() != DialogResult.OK) return;

            txbFilePath.Text = openFileDialog1.FileName;
            if (!File.Exists(txbFilePath.Text))
            {
                Logger.WarnFormat("file not found at {0}", txbFilePath.Text);
                return;
            }

            await Task.Run(() =>
            {
                var peModule = _peService.Load(txbFilePath.Text);
                if (peModule == null) return;

                if (peModule.RelocatedBlocks != null && peModule.RelocatedBlocks.Count > 0)
                {
                    Logger.InfoFormat("Relocations summary: {0} with {1} offsets", peModule.RelocatedBlocks.Count, peModule.RelocatedBlocks.SelectMany(x => x.TypeOffsets).Count(x => x.Type == RelocationType.ImageRelBasedHighlow));
                }
                if (peModule.ImportedModules != null && peModule.ImportedModules.Count > 0)
                {
                    Logger.InfoFormat("Imports summary: {0}", string.Join(", ", peModule.ImportedModules));
                }
                if (peModule.ExportedFunctions != null && peModule.ExportedFunctions.Count > 0)
                {
                    Logger.InfoFormat("Exports summary: {0} function(s)", peModule.ExportedFunctions.Count);
                }

                // view data
                var dosHeaderValues = _dynamicValueAccessor.GetObjectFieldValues(peModule.DosHeader)
                    .Where(x => x.Key.FieldType.IsValueType)
                    .Select(x => new KeyValuePair<string, string>(x.Key.Name, Convert.ToInt32(x.Value).ToString("x8")))
                    .ToList();
                var optionalHeaderValues = _dynamicValueAccessor.GetObjectPropertyValues(peModule.OptionalHeader)
                    .Where(x => x.Key.PropertyType.IsValueType)
                    .Select(x => x.Key.PropertyType.IsPrimitive
                        ? new KeyValuePair<string, string>(x.Key.Name, Convert.ToInt64(x.Value).ToString("x"))
                        : new KeyValuePair<string, string>(x.Key.Name, x.Value.ToString()))
                    .ToList();
                var directories = peModule.OptionalHeader.DataDirectory.Select((dir, index) => new DirectoyViewItem
                {
                    Name = PeDirectories[index],
                    Address = dir.VirtualAddress,
                    Size = dir.Size
                }).ToList();
                var sections = peModule.SectionHeaders.ToList();
                var imports = (from importedModule in peModule.ImportedModules
                    from importedFunction in importedModule.Functions
                    select new ImportedFunctionViewItem
                    {
                        ModuleName = importedModule.ModuleName,
                        Function = importedFunction
                    }).ToList();

                // update view
                Invoke(new Action(() =>
                {
                    ClearTabs();
                    tabControl.TabPages.Add(CreateHeaderTab("DosHeader", dosHeaderValues));
                    tabControl.TabPages.Add(CreateHeaderTab("OptionalHeader", optionalHeaderValues));
                    tabControl.TabPages.Add(CreateDirectoyTab(directories));
                    tabControl.TabPages.Add(CreateSectionTab(sections));
                    if (peModule.ImportedModules != null && peModule.ImportedModules.Count > 0) tabControl.TabPages.Add(CreateImportTab(imports));
                    if (peModule.ExportedFunctions != null && peModule.ExportedFunctions.Count > 0) tabControl.TabPages.Add(CreateExportTab(peModule.ExportedFunctions));
                }));
            }).ConfigureAwait(false);
        }

        private TabPage CreateExportTab(IList<ExportedFunction> dataSource)
        {
            var listViewForm = new ListViewForm<ExportedFunction>(new ListViewFormDataSource<ExportedFunction>
            {
                Columns = new List<ListViewColumn>
                {
                    new ListViewColumn {Name = "Ordinal", Tag = null},
                    new ListViewColumn {Name = "RVA", Tag = null},
                    new ListViewColumn {Name = "Function Name", Tag = null},
                },
                Data = dataSource,
                DataGeters = new List<Func<ExportedFunction, ListViewColumn, string>>
                {
                    (item, column) => item.Ordinal.ToString(),
                    (item, column) => string.Format("0x{0:x8}", item.Rva),
                    (item, column) => item.DecoratedName,
                }
            });
            return CreateTabPage("Exports", listViewForm);
        }

        private TabPage CreateImportTab(IList<ImportedFunctionViewItem> dataSource)
        {
            var listViewForm = new ListViewForm<ImportedFunctionViewItem>(new ListViewFormDataSource<ImportedFunctionViewItem>
            {
                Columns = new List<ListViewColumn>
                {
                    new ListViewColumn {Name = "Module", Tag = null},
                    new ListViewColumn {Name = "Function", Tag = null},
                    new ListViewColumn {Name = "Ordinal", Tag = "number"},
                    new ListViewColumn {Name = "RVA", Tag = null},
                },
                Data = dataSource,
                DataGeters = new List<Func<ImportedFunctionViewItem, ListViewColumn, string>>
                {
                    (item, column) => item.ModuleName,
                    (item, column) => item.Function.DecoratedName,
                    (item, column) => item.Function.Ordinal.ToString(),
                    (item, column) => string.Format("0x{0:x8}", item.Function.ThunkRva),
                }
            });
            return CreateTabPage("Imports", listViewForm);
        }

        public class ImportedFunctionViewItem
        {
            public string ModuleName { get; set; }
            public ImportedFunction Function { get; set; }
        }

        private TabPage CreateSectionTab(IList<SectionHeader> dataSource)
        {
            var listViewForm = new ListViewForm<SectionHeader>(new ListViewFormDataSource<SectionHeader>
            {
                Columns = new List<ListViewColumn>
                {
                    new ListViewColumn {Name = "Name", Tag = null},
                    new ListViewColumn {Name = "VOffset", Tag = null},
                    new ListViewColumn {Name = "VSize", Tag = null},
                    new ListViewColumn {Name = "ROffset", Tag = null},
                    new ListViewColumn {Name = "RSize", Tag = null},
                    new ListViewColumn {Name = "Flags", Tag = null},
                },
                Data = dataSource,
                DataGeters = new List<Func<SectionHeader, ListViewColumn, string>>
                {
                    (item, column) => item.Name,
                    (item, column) => string.Format("0x{0:x8}", item.VirtualAddress),
                    (item, column) => string.Format("0x{0:x8}", item.PhysicalAddressOrVirtualSize),
                    (item, column) => string.Format("0x{0:x8}", item.PointerToRawData),
                    (item, column) => string.Format("0x{0:x8}", item.SizeOfRawData),
                    (item, column) => string.Format("0x{0:x8} {1}", (uint) item.Characteristics, item.Characteristics.ToString()),
                }
            });
            return CreateTabPage("Sections", listViewForm);
        }

        private TabPage CreateDirectoyTab(IList<DirectoyViewItem> dataSource)
        {
            var listViewForm = new ListViewForm<DirectoyViewItem>(new ListViewFormDataSource<DirectoyViewItem>
            {
                Columns = new List<ListViewColumn>
                {
                    new ListViewColumn {Name = "Directory", Tag = null},
                    new ListViewColumn {Name = "RVA", Tag = "number"},
                    new ListViewColumn {Name = "Size", Tag = "number"},
                },
                Data = dataSource,
                DataGeters = new List<Func<DirectoyViewItem, ListViewColumn, string>>
                {
                    (item, column) => item.Name,
                    (item, column) => string.Format("0x{0:x8}", item.Address),
                    (item, column) => string.Format("0x{0:x}", item.Size),
                }
            });
            return CreateTabPage("Directories", listViewForm);
        }

        public class DirectoyViewItem
        {
            public string Name { get; set; }
            public uint Address { get; set; }
            public uint Size { get; set; }
        }

        private TabPage CreateHeaderTab(string tabName, IList<KeyValuePair<string, string>> dataSource)
        {
            var listViewForm = new ListViewForm<KeyValuePair<string, string>>(new ListViewFormDataSource<KeyValuePair<string, string>>
            {
                Columns = new List<ListViewColumn>
                {
                    new ListViewColumn {Name = "Key", Tag = null},
                    new ListViewColumn {Name = "Value", Tag = null},
                },
                Data = dataSource,
                DataGeters = new List<Func<KeyValuePair<string, string>, ListViewColumn, string>>
                {
                    (item, column) => item.Key,
                    (item, column) => item.Value
                }
            });
            return CreateTabPage(tabName, listViewForm);
        }

        private TabPage CreateTabPage(string tabName, Form form)
        {
            // set the form
            form.SetTabStyle().Show();

            // wrap the form in tab page
            var tabPage = new TabPage {Text = tabName};
            tabPage.Controls.Add(form);
            return tabPage;
        }

        private void PeForm_FormClosing(object sender, FormClosingEventArgs e)
        {
            ClearTabs();
        }

        private void ClearTabs()
        {
            // close nested forms
            foreach (TabPage tabPage in tabControl.TabPages)
            {
                var forms = tabPage.Controls.OfType<Form>();
                foreach (var form in forms)
                {
                    form.Close();
                }
                tabPage.Controls.Clear(); // clear controls in tab
            }
            // clear tabs
            tabControl.TabPages.Clear();
        }
    }
}